from stream import Stream
import time

class Valve(Stream):
    """
    Flow control for streams.
    """
    def __init__(self, throughput, interval = 100):
        """
        @param throughput        streamed data size per second.
        @param interval             millisecond
        """
        self.interval = float(interval) / 1000
        self.throughput = throughput * 10
        self.blocksize = int(self.throughput * self.interval)
        self.begintime = time.time()
        self.lasttime = self.begintime
        self.totaldata = 0
        Stream.__init__(self)
        
        #print "interval =", self.interval
        #print "throughput = ", self.throughput
        #print "blocksize", self.blocksize
        
    def write(self, data):
        index = 0
        while index < len(data):
            tail = index + self.blocksize
            if tail > len(data):
                tail = len(data)
            self.doValve(data[index : tail])
            index = tail
            
    def doValve(self, data):
        now = time.time()
        if now - self.lasttime > 3: #no action for 3 seconds
            self.begintime = now
            self.lasttime = now
            self.totaldata = 0
        while self.totaldata > (now - self.begintime) * self.blocksize:
            time.sleep(self.interval)
            now = time.time()
        self.lasttime = now
        for downstream in self.downstream:
            downstream.write(data)
        self.totaldata += len(data)
    
